import json
import os
import getpass
from langchain_core.tools import tool
from langchain.agents import create_agent
from langchain_core.messages import HumanMessage, SystemMessage
from langchain.agents.middleware import HumanInTheLoopMiddleware # 用于人机协作审核
from langgraph.checkpoint.memory import InMemorySaver # 用于保存状态以实现暂停/恢复
from langgraph.types import Command
from langchain_openai import ChatOpenAI
from langchain_core.language_models import BaseChatModel

# 选择一个聊天模型
try:
    # 假设用户已设置必要的API Key
    model = ChatOpenAI(
    model="qwen-plus",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1/",
    api_key="sk-4adf3d4278414c21abf8f26fa9aaa267", # 请替换为您的真实API Key
    temperature=0,
    streaming=True  # 启用流式输出
)
except Exception as e:
    print(f"Warning: Could not initialize model. Please set API key. Error: {e}")
    # 创建一个模拟模型以确保代码结构可运行
    class MockChatModel(BaseChatModel):
        def invoke(self, messages, **kwargs):
            return SystemMessage(content="Mock model response.")
        async def ainvoke(self, messages, **kwargs):
            return SystemMessage(content="Mock model response.")
        @property
        def _llm_type(self) -> str:
            return "mock"
    model = MockChatModel()


# ============================================================================
# 1. 定义底层工具 (Define low-level API tools - Stubbed)
# ============================================================================

@tool
def create_calendar_event(
    title: str,
    start_time: str,  # ISO 格式: "YYYY-MM-DDTHH:MM:SS"
    end_time: str,    # ISO 格式: "YYYY-MM-DDTHH:MM:SS"
    attendees: list[str], # 邮件地址
    location: str = ""
) -> str:
    """创建一个日历事件。需要准确的 ISO 日期时间格式。"""
    return f"Event created: {title} from {start_time} to {end_time} with {len(attendees)} attendees"

@tool
def send_email(
    to: list[str],  # 邮件地址
    subject: str,
    body: str,
    cc: list[str] = []
) -> str:
    """通过邮件 API 发送邮件。需要正确格式的地址。"""
    return f"Email sent to {', '.join(to)} - Subject: {subject}"

@tool
def get_available_time_slots(
    attendees: list[str],
    date: str,  # ISO 格式: "YYYY-MM-DD"
    duration_minutes: int
) -> list[str]:
    """检查指定日期上，参会者的空闲时间段。"""
    return ["09:00", "14:00", "16:00"]

# ============================================================================
# 2. 创建专业化的子代理 (Create specialized sub-agents)
# ============================================================================

# --- 日历代理 ---
CALENDAR_AGENT_PROMPT = (
    "你是一个日历日程安排助理。 "
    "请将自然语言的日程安排请求（例如, '下周二下午 2 点'）解析成正确的 ISO 日期时间格式。 "
    "在需要时使用 get_available_time_slots 检查可用性。 "
    "使用 create_calendar_event 来安排事件。 "
    "在最终回复中，始终确认已安排的内容。"
)

calendar_agent = create_agent(
    model,
    tools=[create_calendar_event, get_available_time_slots],
    system_prompt=CALENDAR_AGENT_PROMPT,
    middleware=[
        HumanInTheLoopMiddleware(
            interrupt_on={"create_calendar_event": True}, 
            description_prefix="Calendar event pending approval",
        ),
    ],
)

# --- 邮件代理 ---
EMAIL_AGENT_PROMPT = (
    "你是一个邮件助理。 "
    "根据自然语言请求撰写专业的邮件。 "
    "提取收件人信息并撰写合适的主题和正文。 "
    "使用 send_email 发送邮件。 "
    "在最终回复中，始终确认已发送的内容。"
)

email_agent = create_agent(
    model,
    tools=[send_email],
    system_prompt=EMAIL_AGENT_PROMPT,
    middleware=[
        HumanInTheLoopMiddleware(
            interrupt_on={"send_email": True}, 
            description_prefix="Outbound email pending approval",
        ),
    ],
)


# ============================================================================
# 3. 将子代理封装为主管代理的工具 (Wrap sub-agents as tools for the supervisor)
# ============================================================================

@tool
def schedule_event(request: str) -> str:
    """使用自然语言安排日历事件。

    当用户想要创建、修改或检查日历约会时使用此工具。
    处理日期/时间解析、可用性检查和事件创建。

    输入: 自然语言日程安排请求 (e.g., '下周二下午 2 点与设计团队开会')
    """
    result = calendar_agent.invoke({
        "messages": [HumanMessage(content=request)]
    })
    return result["messages"][-1].text

@tool
def manage_email(request: str) -> str:
    """使用自然语言发送电子邮件。

    当用户想要发送通知、提醒或任何邮件通信时使用此工具。
    处理收件人提取、主题生成和邮件撰写。

    输入: 自然语言邮件请求 (e.g., '发邮件提醒他们关于会议的事情')
    """
    result = email_agent.invoke({
        "messages": [HumanMessage(content=request)]
    })
    return result["messages"][-1].text


# ============================================================================
# 4. 创建主管代理 (Create the supervisor agent)
# ============================================================================

SUPERVISOR_PROMPT = (
    "你是一个乐于助人的个人助理。 "
    "你可以安排日历事件和发送电子邮件。 "
    "将用户请求分解为适当的工具调用并协调结果。 "
    "当请求涉及多个动作时，按顺序使用多个工具。"
)

supervisor_agent = create_agent(
    model,
    tools=[schedule_event, manage_email],
    system_prompt=SUPERVISOR_PROMPT,
    checkpointer=InMemorySaver(), # 启用状态保存，用于人机审核
)


# ============================================================================
# 5. 使用主管代理 (Use the supervisor)
# ============================================================================

if __name__ == "__main__":
    # 复杂的跨领域请求：安排会议 + 发送邮件提醒
    user_request = (
        "下周二下午 2 点安排与设计团队开会 1 小时，"
        "并给他们发送一封关于回顾新模型图的邮件提醒。"
    )

    print("="*20 + " 初始请求 " + "="*20)
    print(f"用户请求: {user_request}")
    print("="*50 + "\n")

    # 配置：需要一个 thread_id 来保存状态
    config = {"configurable": {"thread_id": "supervisor_run_1"}} 
    
    # 第一次运行：代理会中断等待人机审核
    print("--- 第一次运行 (会中断) ---")
    interrupts = []
    try:
        for step in supervisor_agent.stream(
            {"messages": [HumanMessage(content=user_request)]},
            config,
        ):
            for update in step.values():
                if isinstance(update, dict):
                    # 打印 AI 消息 (包括工具调用)
                    for message in update.get("messages", []):
                        if hasattr(message, 'pretty_print'):
                            message.pretty_print()
                        else:
                            print(message.content)
                else:
                    # 捕获中断事件
                    interrupt_ = update[0]
                    interrupts.append(interrupt_)
                    print(f"\n<<< 发生中断: {interrupt_.id} >>>\n")
    except Exception as e:
        # 捕获流结束时的异常（例如，LangChain-specific异常）
        print(f"流结束或发生异常: {e}")

    # 检查捕获的中断事件
    if not interrupts:
        print("没有发生中断。可能模型决定不调用工具或模型/环境配置问题。")
    else:
        print("="*20 + " 审查请求详情 " + "="*20)
        for i, interrupt_ in enumerate(interrupts):
            for request in interrupt_.value["action_requests"]:
                print(f"中断 ID: {interrupt_.id} ({i+1}/{len(interrupts)})")
                print(f"描述: {request['description']}")
                # 修正: 使用 'name' 键获取工具名称
                tool_name = request.get('name', 'N/A')
                print(f"工具: {tool_name}")
                # 修正: 使用 'args' 键获取参数
                args = request.get('args', {})
                print(f"参数: {json.dumps(args, indent=2, ensure_ascii=False)}\n")

        # --- 模拟人机审核和恢复 ---
        
        # 决定：批准日历事件，并编辑邮件主题后批准发送邮件
        resume_commands = {}
        for interrupt_ in interrupts:
            # 获取待审核的动作请求
            action_request = interrupt_.value["action_requests"][0]
            tool_name = action_request.get('name') # 获取工具名称
            
            if tool_name == "send_email":
                # 编辑邮件：修改主题
                edited_action = action_request.copy()
                # 假设收件人是 design-team@company.com
                edited_action["args"]["to"] = ["design-team@company.com"] # 修正: 使用 "args"
                edited_action["args"]["subject"] = "紧急提醒：新模型图评审" # 修正: 使用 "args"
                resume_commands[interrupt_.id] = {
                    "decisions": [{"type": "edit", "edited_action": edited_action}]
                }
            elif tool_name == "create_calendar_event":
                # 批准日历事件
                resume_commands[interrupt_.id] = {"decisions": [{"type": "approve"}]}
            else:
                # 处理未知的或默认情况，通常是批准
                resume_commands[interrupt_.id] = {"decisions": [{"type": "approve"}]}

        print("="*20 + " 恢复执行 (人机审核决策) " + "="*20)
        print("决策：批准日历事件，编辑并批准邮件发送。")
        print("="*50 + "\n")

        # 恢复执行
        for step in supervisor_agent.stream(
            Command(resume=resume_commands), 
            config,
        ):
            for update in step.values():
                if isinstance(update, dict):
                    for message in update.get("messages", []):
                        if hasattr(message, 'pretty_print'):
                            message.pretty_print()
                        else:
                            print(message.content)